# *-* coding:utf8 *-*

from PyQt5.QtWidgets import QHBoxLayout
from PyQt5.QtWidgets import QVBoxLayout
from PyQt5.QtWidgets import QPushButton
from PyQt5.QtWidgets import QTextEdit
from PyQt5.QtWidgets import QComboBox
from PyQt5.QtWidgets import QLabel
from PyQt5.QtWidgets import QLineEdit
from PyQt5.QtWidgets import QWidget
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtWidgets import QGroupBox
from PyQt5.QtWidgets import QStackedWidget
from PyQt5.QtWidgets import QMessageBox
from PyQt5.QtCore import QRegExp
from PyQt5.QtGui import QRegExpValidator
import traceback
import PHCrypto
import CTFCrypto
import binascii


class CryptoException:
    @staticmethod
    def format_exception_msg(widget, msg):
        QMessageBox.critical(widget, "错误", str(msg), QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
        return 'traceback.format_exc():\n%s' % traceback.format_exc()


class CryptWidget(QWidget):
    crypt_type = 'aes'
    model_type = 'ecb'
    padding_type = 'pkcs5padding'
    data_type = '128'
    out_type = 'base64'
    crypt_list = ['aes', 'des', 'rsa', '3des']
    model_list = ['ecb', 'cbc', 'ofb', 'cfb']
    padding_list = ['pkcs5padding', 'pkcs7padding', 'zeropadding']
    data_list = ['128', '192', '256']
    out_list = ['base64', 'hex']
    padding_dic = {'pkcs7padding': 'pkcs7', 'pkcs5padding': 'pkcs5', 'zeropadding': 'zero'}
    key_text = None
    iv_text = None
    in_text_edit = None
    out_text_edit = None
    stacked_widget = None
    log_out = None
    data_combo = None
    ############################
    # init_rsa_password_widget
    rsa_widget = None
    public_key_text = None
    modulus_text = None
    exponents_text = None
    exponents_or_modulus = '公钥加密'
    rsa_combo = None
    ############################

    def __init__(self, log_out):
        super(CryptWidget, self).__init__()
        self.init_ui()
        self.log_out = log_out

    def init_option_layout(self):
        crypt_label = QLabel('加密类型:')
        crypt_combo = QComboBox()
        crypt_combo.addItems(self.crypt_list)
        crypt_combo.currentTextChanged.connect(self.crypt_change)

        model_label = QLabel('加密模式:')
        model_combo = QComboBox()
        model_combo.addItems(self.model_list)
        model_combo.currentTextChanged.connect(self.model_change)
        model_combo.currentText()

        padding_label = QLabel('填充:')
        padding_combo = QComboBox()
        padding_combo.addItems(self.padding_list)
        padding_combo.currentTextChanged.connect(self.padding_change)

        data_label = QLabel('数据块:')
        self.data_combo = QComboBox()

        self.data_combo.addItems(self.data_list)
        self.data_combo.currentTextChanged.connect(self.data_change)

        out_label = QLabel('输出:')
        out_combo = QComboBox()
        out_combo.addItems(self.out_list)
        out_combo.currentTextChanged.connect(self.out_change)

        h_box_control = QHBoxLayout()
        h_box_control.addWidget(crypt_label)
        h_box_control.addWidget(crypt_combo)

        h_box_control.addWidget(model_label)
        h_box_control.addWidget(model_combo)

        h_box_control.addWidget(padding_label)
        h_box_control.addWidget(padding_combo)

        h_box_control.addWidget(data_label)
        h_box_control.addWidget(self.data_combo)

        h_box_control.addWidget(out_label)
        h_box_control.addWidget(out_combo)

        encrypt_button = QPushButton('加密')
        encrypt_button.clicked.connect(self.encrypt)
        decrypt_button = QPushButton('解密')
        decrypt_button.clicked.connect(self.decrypt)

        import_button = QPushButton('导入')
        import_button.clicked.connect(self.import_texts)
        export_button = QPushButton('导出')
        export_button.clicked.connect(self.export_texts)

        h_box_control.addWidget(encrypt_button)
        h_box_control.addWidget(decrypt_button)

        h_box_control.addWidget(import_button)
        h_box_control.addWidget(export_button)

        return h_box_control

    def init_rsa_password_widget(self):

        self.rsa_combo = QComboBox()
        self.rsa_combo.addItem('公钥加密')
        self.rsa_combo.addItem('模数指数加密')
        self.rsa_combo.currentTextChanged.connect(self.rsa_change)

        h_box_public_key = QHBoxLayout()
        self.public_key_text = QTextEdit()
        h_box_public_key.addWidget(self.public_key_text)
        public_key_group_box = QGroupBox('公钥')
        public_key_group_box.setLayout(h_box_public_key)

        modulus_label = QLabel('指数: ')
        self.modulus_text = QLineEdit()
        h_box_modulus = QHBoxLayout()
        h_box_modulus.addWidget(modulus_label)
        h_box_modulus.addWidget(self.modulus_text)

        h_box_exponents = QHBoxLayout()
        exponents_label = QLabel('模数: ')
        self.exponents_text = QLineEdit()
        h_box_exponents.addWidget(exponents_label)
        h_box_exponents.addWidget(self.exponents_text)

        v_box_exponents_modulus = QVBoxLayout()
        v_box_exponents_modulus.addLayout(h_box_modulus)
        v_box_exponents_modulus.addLayout(h_box_exponents)

        exponents_group_box = QGroupBox()
        exponents_group_box.setLayout(v_box_exponents_modulus)

        v_box = QVBoxLayout()
        v_box.addWidget(exponents_group_box)

        exponents_modulus_widget = QWidget()
        exponents_modulus_widget.setLayout(v_box)

        self.rsa_widget = QStackedWidget()
        self.rsa_widget.addWidget(public_key_group_box)
        self.rsa_widget.addWidget(exponents_modulus_widget)

        h_box = QHBoxLayout()
        h_box.addWidget(self.rsa_combo)
        h_box.addWidget(self.rsa_widget)

        widget = QWidget()
        widget.setLayout(h_box)
        return widget

    def init_password_layout(self):

        key_label = QLabel('key:')
        self.key_text = QLineEdit()
        self.key_text.setMaxLength(16)
        #self.key_text.setValidator(QRegExpValidator(QRegExp("{16}"), self))
        h_box_key = QHBoxLayout()
        h_box_key.addWidget(key_label)
        h_box_key.addWidget(self.key_text)

        iv_label = QLabel('i v:')
        self.iv_text = QLineEdit()
        self.iv_text.setEnabled(False)
        self.iv_text.setMaxLength(16)
        #self.iv_text.setValidator(QRegExpValidator(QRegExp("{16}"), self))
        h_box_iv = QHBoxLayout()
        h_box_iv.addWidget(iv_label)
        h_box_iv.addWidget(self.iv_text)

        v_box_key_or_iv = QVBoxLayout()
        v_box_key_or_iv.addLayout(h_box_key)
        v_box_key_or_iv.addLayout(h_box_iv)

        widget = QWidget()
        widget.setLayout(v_box_key_or_iv)

        self.stacked_widget = QStackedWidget()
        self.stacked_widget.addWidget(widget)
        self.stacked_widget.addWidget(self.init_rsa_password_widget())

        h_box_password = QHBoxLayout()
        h_box_password.addWidget(self.stacked_widget)
        return h_box_password

    def init_in_out_layout(self):
        in_group_box = QGroupBox('输入')
        in_layout = QHBoxLayout()
        self.in_text_edit = QTextEdit()
        self.in_text_edit.setFontFamily('Consolas')
        self.in_text_edit.setFontPointSize(14)
        in_layout.addWidget(self.in_text_edit)
        in_group_box.setLayout(in_layout)

        out_group_box = QGroupBox('输出')
        out_layout = QHBoxLayout()
        self.out_text_edit = QTextEdit()
        self.out_text_edit.setFontFamily('Consolas')
        self.out_text_edit.setFontPointSize(14)
        self.out_text_edit.setReadOnly(True)
        out_layout.addWidget(self.out_text_edit)
        out_group_box.setLayout(out_layout)

        v_box_text = QVBoxLayout()
        v_box_text.stretch(2)
        v_box_text.addWidget(in_group_box, 1)
        v_box_text.addWidget(out_group_box, 1)

        return v_box_text

    def init_ui(self):
        v_box = QVBoxLayout()
        v_box.stretch(100)
        v_box.addLayout(self.init_option_layout(), 1)
        v_box.addLayout(self.init_password_layout(), 1)
        v_box.addLayout(self.init_in_out_layout(), 98)
        self.setLayout(v_box)

    def get_param(self):
        padding = self.padding_dic[self.padding_type]
        key = self.key_text.text()
        iv = self.iv_text.text()
        return padding, key, iv

    def encrypt(self):
        try:
            padding, key, iv = self.get_param()
            self.out_text_edit.clear()
            if self.crypt_type == 'aes':
                if self.model_type == 'ecb':
                    aes = PHCrypto.AESCrypto(key, padding)  # d8cg8gVakEq9Agup
                else:
                    aes = PHCrypto.AESCrypto(key, padding, iv)
                self.batch_encrypt(aes)
            elif self.crypt_type == 'des':
                if self.model_type == 'ecb':
                    des = PHCrypto.DESCrypto(key, padding)
                else:
                    des = PHCrypto.DESCrypto(key, padding, iv)
                self.batch_encrypt(des)
            elif self.crypt_type == '3des':
                if self.model_type == 'ecb':
                    des = PHCrypto.DES3Crypto(key, padding)
                else:
                    des = PHCrypto.DES3Crypto(key, padding, iv)
                self.batch_encrypt(des)
            elif self.crypt_type == 'rsa':
                self.rsa_encrypt()
        except Exception as e:
            self.log_out.append(CryptoException.format_exception_msg(self, e))

    def decrypt(self):
        try:
            padding, key, iv = self.get_param()
            self.out_text_edit.clear()
            if self.crypt_type == 'aes':
                if self.model_type == 'ecb':
                    aes = PHCrypto.AESCrypto(key, padding)  # d8cg8gVakEq9Agup
                else:
                    aes = PHCrypto.AESCrypto(key, padding, iv)
                self.batch_decrypt(aes)
            elif self.crypt_type == 'des':
                if self.model_type == 'ecb':
                    des = PHCrypto.DESCrypto(key, padding)
                else:
                    des = PHCrypto.DESCrypto(key, padding, iv)
                self.batch_decrypt(des)
            elif self.crypt_type == '3des':
                if self.model_type == 'ecb':
                    des = PHCrypto.DES3Crypto(key, padding)
                else:
                    des = PHCrypto.DES3Crypto(key, padding, iv)
                self.batch_decrypt(des)
        except Exception as e:
            self.log_out.append(CryptoException.format_exception_msg(self, e))

    # 批量加密
    def batch_encrypt(self, crypto):
        texts = self.get_texts(self.in_text_edit)
        for text in texts:
            try:
                if self.model_type == 'ecb':
                    encrypt_text = crypto.ecb_encrypt(text)
                elif self.model_type == 'cbc':
                    encrypt_text = crypto.cbc_encrypt(text)
                elif self.model_type == 'ofb':
                    encrypt_text = crypto.ofb_encrypt(text)
                elif self.model_type == 'cfb':
                    encrypt_text = crypto.cfb_encrypt(text)

                if self.out_type == 'hex':
                    encrypt_text = self.base64_to_hex(encrypt_text)
                self.out_text_edit.append(encrypt_text)
            except Exception as e:
                self.log_out.append(CryptoException.format_exception_msg(self, e))

    # 批量解密
    def batch_decrypt(self, crypto):
        texts = self.get_texts(self.in_text_edit)
        for text in texts:
            try:
                if self.model_type == 'ecb':
                    encrypt_text = crypto.ecb_decrypt(text)
                elif self.model_type == 'cbc':
                    encrypt_text = crypto.cbc_decrypt(text)
                elif self.model_type == 'ctr':
                    encrypt_text = crypto.ctr_decrypt(text)
                elif self.model_type == 'ofb':
                    encrypt_text = crypto.ofb_decrypt(text)
                elif self.model_type == 'cfb':
                    encrypt_text = crypto.cfb_decrypt(text)

                if self.out_type == 'hex':
                    encrypt_text = self.base64_to_hex(encrypt_text)
                self.out_text_edit.append(encrypt_text)
            except Exception as e:
                self.log_out.append(CryptoException.format_exception_msg(self, e))

    def rsa_encrypt(self):
        rsa_crypto = None
        if self.exponents_or_modulus == '公钥加密':
            public_key_text = self.public_key_text.toPlainText()
            rsa_crypto = PHCrypto.RSACrypto(public_key_text)

        elif self.exponents_or_modulus == '模数指数加密':
            modulus = self.modulus_text.text()
            exponents = self.exponents_text.text()
            rsa_crypto = PHCrypto.RSACrypto('', modulus, exponents)

        self.out_text_edit.clear()
        texts = self.get_texts(self.in_text_edit)
        for text in texts:
            try:
                encrypt_text = rsa_crypto.encrypt(text)
                if self.out_type == 'hex':
                    encrypt_text = self.base64_to_hex(rsa_crypto)
                self.out_text_edit.append(encrypt_text)
            except Exception as e:
                self.log_out.append(CryptoException.format_exception_msg(self, e))

    def out_text(self):
        texts = self.get_texts(self.out_text_edit)
        self.out_text_edit.clear()

        if self.out_type == 'base64':
            try:
                for text in texts:
                    base = self.hex_to_base64(text)
                    self.out_text_edit.append(base)
            except Exception as e:
                print(e)
        elif self.out_type == 'hex':
            try:
                for text in texts:
                    hex_str = self.base64_to_hex(text)
                    self.out_text_edit.append(hex_str)
            except Exception as e:
                print(e)

    def crypt_change(self, text):
        if text == 'rsa':
            self.stacked_widget.setCurrentIndex(1)
            self.rsa_combo.setEnabled(True)
        elif text == 'sm2':
            self.rsa_combo.setEnabled(False)
            self.stacked_widget.setCurrentIndex(1)
        elif text == 'des':
            self.data_combo.setEnabled(False)
            self.key_text.setMaxLength(8)
            # self.key_text.setValidator(QRegExpValidator(QRegExp("[0-9a-zA-Z]{8}"), self))
            if len(self.key_text.text()) > 8:
                self.key_text.setText(self.key_text.text()[:8])
            self.stacked_widget.setCurrentIndex(0)
            self.iv_text.setMaxLength(8)
        elif text == '3des':
            self.data_combo.setEnabled(False)
            self.stacked_widget.setCurrentIndex(0)
        else:
            self.stacked_widget.setCurrentIndex(0)
            self.key_text.setMaxLength(16)
            # self.key_text.setValidator(QRegExpValidator(QRegExp("[0-9a-zA-Z]{16}"), self))
            self.data_combo.setEnabled(True)
            self.iv_text.setMaxLength(16)
        self.crypt_type = text

    def model_change(self, text):
        if text == 'ecb':
            self.iv_text.setEnabled(False)
            self.iv_text.clear()
        else:
            self.iv_text.setEnabled(True)

        self.model_type = text

    def padding_change(self, text):
        self.padding_type = text

    def data_change(self, text):
        if '128' == text:
            self.key_text.setMaxLength(16)
            # self.key_text.setValidator(QRegExpValidator(QRegExp("{16}"), self))
            if len(self.key_text.text()) > 16:
                self.key_text.setText(self.key_text.text()[:16])
        elif '192' == text:
            self.key_text.setMaxLength(24)
            # self.key_text.setValidator(QRegExpValidator(QRegExp("[0-9a-zA-Z]{24}"), self))
            if len(self.key_text.text()) > 24:
                self.key_text.setText(self.key_text.text()[:24])
        elif '256' == text:
            self.key_text.setMaxLength(32)
            # self.key_text.setValidator(QRegExpValidator(QRegExp("[0-9a-zA-Z]{32}"), self))

        self.data_type = text

    def out_change(self, text):
        self.out_type = text
        self.out_text()

    def rsa_change(self, text):
        if text == '公钥加密':
            self.rsa_widget.setCurrentIndex(0)
        elif text == '模数指数加密':
            self.rsa_widget.setCurrentIndex(1)
        self.exponents_or_modulus = text

    def import_texts(self):
        file_name, file_type = QFileDialog.getOpenFileName(self, "导入文件")

        if file_name == "":
            return

        file = open(file_name, 'r')
        lines = file.readlines()
        file.close()
        for line in lines:
            self.in_text_edit.append(line.replace('\n', ''))

    def export_texts(self):
        file_name, file_type = QFileDialog.getSaveFileName(self, "保存文件")

        if file_name == "":
            return
        lines = self.out_text_edit.toPlainText()
        file = open(file_name, 'w')
        file.writelines(lines)
        file.close()

    @staticmethod
    def base64_to_hex(text):
        return text.encode().hex()

    @staticmethod
    def hex_to_base64(hex_str):
        return binascii.a2b_hex(hex_str).decode()

    @staticmethod
    def get_texts(edit):
        texts = edit.toPlainText()
        texts = texts.split('\n')
        texts = [i for i in texts if i != '']
        return texts


class CTFWidget(QWidget):
    in_text_edit = None
    out_text_edit = None
    crypto_combo_box = None
    crypto_list = ['十六进制转换', '栅栏密码', '摩斯密码', 'url编码', '云影密码', '凯撒密码']

    def __init__(self, log_out):
        super(CTFWidget, self).__init__()
        self.log_out = log_out
        self.init_ui()

    def init_in_out_layout(self):
        in_group_box = QGroupBox('输入')
        in_layout = QHBoxLayout()
        self.in_text_edit = QTextEdit()
        self.in_text_edit.setFontFamily('Consolas')
        self.in_text_edit.setFontPointSize(14)
        in_layout.addWidget(self.in_text_edit)
        in_group_box.setLayout(in_layout)

        out_group_box = QGroupBox('输出')
        out_layout = QHBoxLayout()
        self.out_text_edit = QTextEdit()
        self.out_text_edit.setFontFamily('Consolas')
        self.out_text_edit.setFontPointSize(14)
        self.out_text_edit.setReadOnly(True)
        out_layout.addWidget(self.out_text_edit)
        out_group_box.setLayout(out_layout)

        v_box_text = QVBoxLayout()
        v_box_text.stretch(2)
        v_box_text.addWidget(in_group_box, 1)
        v_box_text.addWidget(out_group_box, 1)

        return v_box_text

    def init_ui(self):
        self.crypto_combo_box = QComboBox()
        self.crypto_combo_box.addItems(self.crypto_list)

        encode_button = QPushButton('编码')
        encode_button.clicked.connect(self.encode)

        decode_button = QPushButton('解码')
        decode_button.clicked.connect(self.decode)

        h_box = QHBoxLayout()
        h_box.addWidget(self.crypto_combo_box)
        h_box.addWidget(encode_button)
        h_box.addWidget(decode_button)

        v_box = QVBoxLayout()
        v_box.addLayout(h_box)
        v_box.addLayout(self.init_in_out_layout())
        self.setLayout(v_box)

    def batch_decode(self, crypto):
        texts = self.get_texts(self.in_text_edit)
        self.out_text_edit.clear()
        for text in texts:
            try:
                result = crypto.decode(text)
                self.out_text_edit.append(result)
            except Exception as e:
                self.log_out.append(CryptoException.format_exception_msg(self, e))

    def batch_encode(self, crypto):
        texts = self.get_texts(self.in_text_edit)
        self.out_text_edit.clear()
        for text in texts:
            try:
                result = str(crypto.encode(text))
                self.out_text_edit.append(result)
            except Exception as e:
                self.log_out.append(CryptoException.format_exception_msg(self, e))

    def decode(self):
        ret = self.choose_crypto()
        if ret:
            self.batch_decode(ret)

    def encode(self):
        ret = self.choose_crypto()
        if ret:
            self.batch_encode(ret)

    def choose_crypto(self):
        text = self.crypto_combo_box.currentText()
        crypto = None
        if '十六进制转换' == text:
            crypto = CTFCrypto.HexToString()
        elif '栅栏密码' == text:
            crypto = CTFCrypto.Fence(4)
        elif '摩斯密码' == text:
            crypto = CTFCrypto.Morse()
        elif 'url编码' == text:
            crypto = CTFCrypto.UrlDecode()
        elif '云影密码' == text:
            crypto = CTFCrypto.CloudShadow()
        elif '凯撒密码' == text:
            crypto = CTFCrypto.Caesar(4)
        return crypto

    @staticmethod
    def get_texts(edit):
        texts = edit.toPlainText()
        texts = texts.split('\n')
        texts = [i for i in texts if i != '']
        return texts
